import os
from qagen.utils import *
from random import *
import bpy
import mathutils
from mathutils import Euler, Matrix, Vector
import math
import shutil
from PIL import Image
import numpy as np
from concurrent.futures import ThreadPoolExecutor, as_completed
from queue import Queue
import multiprocessing

cate_list = read_category_list()
obj_list = read_obj_info()
color_info = read_color_info()

file_prefix = '~/scenegen'
data_folder = os.path.expanduser(f'{file_prefix}/data_extra3')

def get_table():
    table_obj = None
    for obj in bpy.data.objects:
        if "TableDining" in obj.name:
            table_obj = obj
            break
    return table_obj
def get_table_():
    table_obj = None
    for obj in bpy.data.objects:
        if "TableDining" in obj.name:
            table_obj = obj
    return table_obj

def get_lamp():
    table_obj = None
    for obj in bpy.data.objects:
        if "PointLamp" in obj.name:
            table_obj = obj
            break
    return table_obj

def get_object_world_bbox(obj):
    mesh = obj.data
    world_verts = [obj.matrix_world @ v.co for v in mesh.vertices]
    # print(len(mesh.vertices))
    return world_verts

def getBbox(parent_obj):
    all_points = []
    bpy.context.view_layer.update()
    if 'TableDining' in parent_obj.name:
        mesh = parent_obj.data
        worlds_verts = [Euler(parent_obj.rotation_euler, parent_obj.rotation_mode).to_matrix().to_4x4() @ v.co for v in mesh.vertices]
        all_points.extend(worlds_verts)
    for obj in parent_obj.children_recursive:
        if obj.type == 'MESH' and not obj.hide_viewport and not obj.hide_render:
            part_bbox = get_object_world_bbox(obj)
            all_points.extend(part_bbox)
    min_corner = Vector((
        min(p.x for p in all_points),
        min(p.y for p in all_points),
        min(p.z for p in all_points)
    ))
    max_corner = Vector((
        max(p.x for p in all_points),
        max(p.y for p in all_points),
        max(p.z for p in all_points)
    ))

    return min_corner, max_corner

def fixLocation(obj, current_center, target_center):
    # global depsgraph
    for obj in obj.children_recursive:
        if obj.children == ():
            obj.location = obj.location - current_center + target_center
            # depsgraph = bpy.context.evaluated_depsgraph_get()

def random_position(sz, tablemin, tablemax):
    tx_min = tablemin.x + sz.x / 2
    tx_max = tablemax.x - sz.x / 2
    ty_min = tablemin.y + sz.y / 2
    ty_max = tablemax.y - sz.y / 2
    tx = uniform(tx_min, tx_max)
    ty = uniform(ty_min, ty_max)
    tz = tablemax.z
    return (tx, ty, tz)

ITV = 0.15
def checkOverlap(obj_nbbox, x, y, sz):
    bxmin = x - sz.x / 2
    bxmax = x + sz.x / 2
    bymin = y - sz.y / 2
    bymax = y + sz.y / 2
    for a_min, a_max in obj_nbbox:
        axmin, aymin, _ = a_min
        axmax, aymax, _ = a_max
        if axmin >= bxmax + ITV or axmax + ITV <= bxmin or aymin >= bymax + ITV or aymax + ITV <= bymin:
            continue
        return True
    return False

def checkBehind(x, y, stx, sty, edx, edy):
    #   (edx - stx, edy - sty) cross times (x - stx, y - sty)
    # = (edx - stx) * (y - sty) - (edy - sty) * (x - stx)
    cross_product = (edx - stx) * (y - sty) - (edy - sty) * (x - stx)
    return cross_product > 0

def recursive_select(obj):
    obj.select_set(True)
    for child in obj.children_recursive:
        child.select_set(True)

def delete_sketchfab_models(blend_file_path):
    bpy.ops.object.select_all(action='DESELECT')
    for obj in bpy.data.objects:
        if "Sketchfab_model" in obj.name:
            recursive_select(obj)
    bpy.ops.object.delete()
    bpy.ops.wm.save_mainfile(filepath=blend_file_path)

def delete_pillars(blend_file_path):
    bpy.ops.object.select_all(action='DESELECT')
    for obj in bpy.data.objects:
        if "PillarFactory" in obj.name:
            recursive_select(obj)
    bpy.ops.object.delete()
    bpy.ops.wm.save_mainfile(filepath=blend_file_path)

def delete_object(obj):
    bpy.ops.object.select_all(action='DESELECT')
    recursive_select(obj)
    bpy.ops.object.delete()

def fix_table():
    bpy.ops.object.select_all(action='DESELECT')
    table_obj = get_table_()
    table_obj.select_set(True)
    table_obj.rotation_euler.x = 0
    table_obj.rotation_euler.y = 0
    table_obj = get_table()
    table_obj.select_set(True)
    table_obj.rotation_euler.x = 0
    table_obj.rotation_euler.y = 0

def random_camera():
    table_obj = get_table()
    if table_obj is None:
        print("未找到名称包含 'TableDining' 的物体！")
    else:
        print(table_obj.location)
        min_corner, max_corner = getBbox(table_obj)
        print(min_corner, max_corner)
        center = (min_corner + max_corner) / 2.0 + table_obj.location # HERE
        width = max_corner.x - min_corner.x
        depth = max_corner.y - min_corner.y
        top_z = max_corner.z

        camera_obj = None
        camrig = None
        for obj in bpy.data.objects:
            if obj.type == 'CAMERA':
                camera_obj = obj
                # print(obj.name, obj.type, obj.data.angle_x, obj.data.angle_y)
            if obj.name == 'camrig.0':
                camrig = obj
                # print(obj.location, obj.rotation_euler)

        cam_data = camera_obj.data

        fov_x = cam_data.angle_x
        fov_y = cam_data.angle_y

        req_distance_x = (width / 2.0) / math.tan(fov_x / 2.0)
        req_distance_y = (depth / 2.0) / math.tan(fov_y / 2.0)
        req_distance = max(req_distance_x, req_distance_y)
        print("计算所需的距离: %.3f" % req_distance)

        rand_ang = randint(45, 50)
        # req_distance += math.log2(width / depth) * 0.3
        req_distance += (width - depth) * depth * 0.25
        camrig.location = (center.x + req_distance * math.sin(math.radians(rand_ang)), center.y, top_z + req_distance - 0.1)
        camrig.rotation_euler = (math.radians(rand_ang), 0, math.radians(90))
        print("摄像机新位置:", camrig.location)
        print("摄像机新旋转 (度):",
              (math.degrees(camrig.rotation_euler.x),
               math.degrees(camrig.rotation_euler.y),
               math.degrees(camrig.rotation_euler.z)))

        # bpy.ops.wm.save_as_mainfile(filepath=scene_path)
        print(f'return {camrig.location=}, {camrig.rotation_euler=}, {fov_x=}, {fov_y=}')
        return camrig.location, camrig.rotation_euler, fov_x, fov_y

def random_len(scl, factor = uniform(0.9, 1.2)):
    scl = max(scl, 10)
    scl = min(scl, 50)
    scl = scl * 0.01 * factor
    return scl

def sample_x_coordinate(index, total, table_min, table_max, sz):
    if total <= 1:
        weight = random()
    else:
        # beta_param = 4 - 3 * (index / (total - 1))
        weight = betavariate(index + 1, total - index)
    x = table_min.x + sz.x / 2 + weight * (table_max.x - table_min.x - sz.x)
    return x

def sample_y_coordinate(table_min, table_max, sz, dir):
    if dir.y == 1:
        return table_min.y + sz.y / 2 + random() * (table_max.y - table_min.y - sz.y) / 2
    elif dir.y == -1:
        return table_max.y - sz.y / 2 - random() * (table_max.y - table_min.y - sz.y) / 2
    return table_min.y + sz.y / 2 + random() * (table_max.y - table_min.y - sz.y)

def getColorId(co):
    if co == 'orange' or co == 'red' or co == 'salmon' or co == 'pink' or co == 'orange-red' or co == 'red-orange':
        return 0
    if co == 'black' or co == 'gray':
        return 1
    if co == 'blue' or co == 'cyan' or co == 'lightblue':
        return 2
    if co == 'green' or co == 'lime' or co == 'lightgreen':
        return 3
    if co == 'yellow' or co == 'gold' or co == 'lightyellow':
        return 4
    if co == 'purple' or co == 'violet' or co == 'lavender':
        return 5
    if co == 'white' or co == 'beige' or co == 'lightgray':
        return 6
    if co == 'brown' or co == 'tan' or co == 'khaki':
        return 7
    return None
def checkColor(color1, color2):
    cid1, cid2 = getColorId(color1), getColorId(color2)
    if cid1 == None or cid2 == None:
        return True
    if cid1 != cid2:
        return True
    return False

def save_images(source_image, source_depth, target_image, target_depth):
    with Image.open(source_image) as img:
        if img.mode in ('RGBA', 'LA'):
            img = img.convert('RGB')
        img.save(target_image, 'JPEG')
        data = np.load(source_depth)
        data_min = np.min(data)
        data_max = min(np.max(data), 4)
        data = 255 - 255 * (np.minimum(data, 4) - data_min) / (data_max - data_min)
        data = np.stack((data,) * 3, axis=-1).astype(np.uint8)
        data = Image.fromarray(data)
        data.save(target_depth, 'PNG')

def checkOverlap_(obj_info, objid):
    for idx, (obj_id, info) in enumerate(obj_info.items()):
        if objid == obj_id:
            continue
        if obj_info[objid]['bboxmin'][0] >= info['bboxmax'][0] or obj_info[objid]['bboxmin'][1] >= info['bboxmax'][1] or info['bboxmin'][0] >= obj_info[objid]['bboxmax'][0] or obj_info[objid]['bboxmax'][1] <= info['bboxmin'][1]:
            continue
        return True
    return False

def genqa(id, scene_id):
    # copy folder

    try:
        target_folder = os.path.expanduser(f'{file_prefix}/scene_{id}')
        shutil.copytree(get_scene_folder(scene_id), target_folder)
        scene_path = get_scene_path_(target_folder)

        bpy.ops.wm.open_mainfile(filepath=scene_path)

        # delete_sketchfab_models(scene_path)
        # delete_pillars(scene_path)
        fix_table()

        cam_loc, cam_rot, cam_x, cam_y = random_camera()
        cam_loc = list(cam_loc)
        cam_rot = list(cam_rot)
        bpy.ops.wm.save_as_mainfile()

        table = get_table()
        table_min, table_max = getBbox(table)
        table_sz = table_max - table_min
        table_min += table.location
        table_max += table.location
        print(f'{table_min=}, {table_max=}')
        table_max.x -= table_sz.x * 0.20
        table_min.y += table_sz.y * 0.10
        table_max.y -= table_sz.y * 0.10

        while True:
            objs_id = sample(list(obj_list.keys()), k = randint(3, 4))
            error = False
            for id1 in objs_id:
                for id2 in objs_id:
                    if id1 != id2 and obj_list[id1]['cate'] == obj_list[id2]['cate']:
                        error = True
            if error:
                continue

            bpy.ops.wm.open_mainfile(filepath=scene_path)
            obj_info = {}
            cnt = 0
            for obj_id in objs_id:
                info = obj_list[obj_id].copy()
                info['scale'] = random_len(info['scale'])
                bpy.ops.import_scene.gltf(filepath=f'myobj/{info["path"]}')
                tmp_list = bpy.context.selected_objects

                armature = None
                for obj in tmp_list:
                    if obj.type == 'ARMATURE':
                        armature = obj
                        break
                if armature != None:
                    obj_info = None
                    break

                obj = tmp_list[0]
                info['object'] = obj
                old_dir = info['direction']
                info['direction'] = mathutils.Vector((old_dir[1], -old_dir[0], old_dir[2]))

                # set scale
                omin, omax = getBbox(obj)
                # print(f'original bbox: {omax - omin}')
                obj.scale *= info['scale'] / max(omax - omin)
                bpy.context.view_layer.update()

                # set rotation
                if abs(old_dir[2]) == 1:
                    ang = uniform(0, 2 * math.pi)
                    # obj.rotation_euler[2] = ang
                    q = mathutils.Quaternion((0, 0, 1), ang)
                    obj.rotation_quaternion = q @ obj.rotation_quaternion
                else:
                    omin, omax = getBbox(obj)
                    from_dir = info['direction']
                    sz = omax - omin
                    if (omax.y - omin.y < 0.05 and from_dir.x != 0) or (omax.x - omin.x < 0.05 and from_dir.y != 0):
                        target = choice([mathutils.Vector((0, 1, 0)), mathutils.Vector((0, -1, 0))])
                    if (omax.x - omin.x < 0.05 and from_dir.x != 0) or (omax.y - omin.y < 0.05 and from_dir.y != 0):
                        target = mathutils.Vector((1, 0, 0))
                    else:
                        target = choice([mathutils.Vector((1, 0, 0)), mathutils.Vector((0, 1, 0)), mathutils.Vector((0, -1, 0))])
                    # print(target)
                    angle_from = math.atan2(from_dir.y, from_dir.x)
                    angle_to = math.atan2(target.y, target.x)
                    # print(obj.rotation_mode)
                    # obj.rotation_euler[2] = angle_to - angle_from
                    q = mathutils.Quaternion((0, 0, 1), angle_to - angle_from)
                    obj.rotation_quaternion = q @ obj.rotation_quaternion
                    info['direction'] = target
                bpy.context.view_layer.update()

                info['bboxmin'], info['bboxmax'] = getBbox(obj)
                cnt += 1
                obj_info[f'{obj_id}_{cnt}'] = info
            if obj_info == None:
                continue
            # obj_info = dict(sorted(obj_info.items(), key=lambda item: item[1]['bboxmax'].z - item[1]['bboxmin'].z, reverse=True))
            
            # input()
            try_times = 0
            while try_times < 20:
                try_times += 1
                obj_nbbox = []
                cnt = 0
                for idx, (obj_id, info) in enumerate(obj_info.items()):
                    obj = info['object']
                    obj_sz = info['bboxmax'] - info['bboxmin']
                    try_times_x = 0
                    while try_times_x < 20:
                        try_times_x += 1
                        x = sample_x_coordinate(idx, len(obj_info), table_min, table_max, obj_sz)
                        try_times_y = 0
                        while try_times_y < 20:
                            try_times_y += 1
                            y = sample_y_coordinate(table_min, table_max, obj_sz, info['direction'])
                            if not checkOverlap(obj_nbbox, x, y, obj_sz):
                                break
                            y = None
                        if y != None:
                            break
                        x = None
                    if x == None:
                        obj_nbbox = None
                        break
                    target_co = mathutils.Vector((x, y, table_max.z + obj_sz.z / 2))
                    obj_nbbox.append((target_co - obj_sz / 2, target_co + obj_sz / 2))
                    source_co = (info['bboxmax'] + info['bboxmin']) / 2
                    obj.location = obj.location - source_co + target_co
                    bpy.context.view_layer.update()
                    info['bboxmin'], info['bboxmax'] = getBbox(obj)
                    obj_nbbox.append((info['bboxmin'], info['bboxmax']))
                if obj_nbbox != None:
                    break
            if obj_nbbox != None:
                break
        if obj_nbbox == None:
            assert(False)
        bpy.ops.wm.save_as_mainfile()

        render_(target_folder)
        save_images(get_image_path_(target_folder), get_depth_path_(target_folder), os.path.expanduser(f'{data_folder}/{id}.jpg'), os.path.expanduser(f'{data_folder}/{id}.png'))

        with open(get_objid_path_(target_folder), 'r') as f:
            id_json = json.load(f)
        # read mask.npy
        mask = np.load(get_mask_path_(target_folder))
        allinfo = {}
        allinfo['type'] = 'timeseq'
        allinfo['scene_id'] = scene_id
        allinfo['camera'] = {
            'location': cam_loc,
            'rotation': cam_rot,
            'fov': [cam_x, cam_y],
        }
        allinfo['table'] = {
            'tablemin': [table_min.x, table_min.y, table_min.z],
            'tablemax': [table_max.x, table_max.y, table_max.z],
        }
        for obj_id, info in obj_info.items():
            info['mask'] = get_object_mask_rle(id_json, mask, info['object'].name)
            # info['object'] = None
            info['bboxmin'] = [info['bboxmin'].x, info['bboxmin'].y, info['bboxmin'].z]
            info['bboxmax'] = [info['bboxmax'].x, info['bboxmax'].y, info['bboxmax'].z]
            info['direction'] = [info['direction'].x, info['direction'].y, info['direction'].z]
            obj_info[obj_id] = info
        allinfo['objects'] = obj_info

        for obj_id, info in obj_info.items():
            zt = info['object'].location.z
            info['object'].location.z = 100000
            bpy.context.view_layer.update()
            bpy.ops.wm.save_as_mainfile()
            render_(target_folder)
            save_images(get_image_path_(target_folder), get_depth_path_(target_folder), os.path.expanduser(f'{data_folder}/{id}_{obj_id[-1]}.jpg'), os.path.expanduser(f'{data_folder}/{id}_{obj_id[-1]}.png'))
            info['object'].location.z = zt
            bpy.context.view_layer.update()
            bpy.ops.wm.save_as_mainfile()
        
        for obj_id, info in obj_info.items():
            # print(table_min, table_max, info['bboxmin'], info['bboxmax'])
            # input()
            backup_loc = info['object'].location.copy()
            times = 0
            while times < 20:
                times += 1
                dir = choice([(0, 1), (1, 0), (0, -1), (-1, 0), (1, 1), (1, -1), (-1, 1), (-1, -1)])
                dist = uniform(0.10, 0.60)
                delta = (dir[0] * dist, dir[1] * dist)
                # print(delta)
                info['object'].location.x = backup_loc.x + delta[0]
                info['object'].location.y = backup_loc.y + delta[1]
                bpy.context.view_layer.update()
                bpy.ops.wm.save_as_mainfile()
                info['bboxmin'], info['bboxmax'] = getBbox(info['object'])
                # print(table_min, table_max, info['bboxmin'], info['bboxmax'])
                # input()
                if info['bboxmax'].x < table_min.x or info['bboxmax'].y < table_min.y or info['bboxmin'].x > table_max.x or info['bboxmin'].y > table_max.y:
                    continue
                if checkOverlap_(obj_info, obj_id):
                    continue
                info['move'] = {
                    'dir': dir,
                    'dist': dist
                }
                render_(target_folder)
                save_images(get_image_path_(target_folder), get_depth_path_(target_folder), os.path.expanduser(f'{data_folder}/{id}_{obj_id[-1]}_.jpg'), os.path.expanduser(f'{data_folder}/{id}_{obj_id[-1]}_.png'))
                mask = np.load(get_mask_path_(target_folder))
                with open(get_objid_path_(target_folder), 'r') as f:
                    id_json = json.load(f)
                info['mask'] = get_object_mask_rle(id_json, mask, info['object'].name)
                break
            info['object'].location = backup_loc
            bpy.context.view_layer.update()
            bpy.ops.wm.save_as_mainfile()
            info['bboxmin'], info['bboxmax'] = getBbox(obj)

        for obj_id, info in obj_info.items():
            info['object'] = None
            info['bboxmin'] = list(info['bboxmin'])
            info['bboxmax'] = list(info['bboxmax'])
        with open(os.path.expanduser(f'{data_folder}/{id}.json'), 'w') as f:
            json.dump(allinfo, f, indent=4)
    
        # delete folder
        shutil.rmtree(target_folder)
    except:
        shutil.rmtree(target_folder, ignore_errors=True)
        return

if __name__ == "__main__":
    os.makedirs(os.path.expanduser(f'{data_folder}'), exist_ok=True)

    with open('Xscene_ids.json', 'r') as f:
        scene_ids = json.load(f)
    
    num_processes = 4
    num_samples = 1024
    # num_samples = 4
    # seed(2333)
    pool = multiprocessing.Pool(processes=num_processes)
    results = []

    for i in range(num_samples):
        id = random_hex()
        scene_id_to_use = choice(scene_ids)
        # scene_id_to_use = scene_ids[i % len(scene_ids)] # 循环使用预选的场景
        results.append(pool.apply_async(genqa, args=(id, scene_id_to_use)))

    # 等待所有进程完成
    for res in results:
        res.get()

    pool.close()
    pool.join()
